#-*- coding:utf-8 -*-
import tensorflow as tf
from getStrokeLabel import genStrokes
#from getStroke import genStrokes
import numpy as np
from config import Config
from myCNN import CNN2
import os
import h5py
import matplotlib.pyplot as plt
import cv2


def read_data(path):
     images = np.array(h5py.File(path, 'r').get('images'))
     labels = np.array(h5py.File(path, 'r').get('labels'))
     return images, labels


config = Config()
draw = genStrokes(config, )
cnn2 = CNN2(config,)


def MSE(image,gt):
      if len(image.shape)>2:
           image=image[:,:,0]
      if len(gt.shape)>2:
           gt=gt[:,:,0]
  
      width=image.shape[0]
      height=image.shape[1]
      sumup=0
      for j in range(width):
         for k in range(height):
             sumup=sumup+(image[j,k]-gt[j,k])*(image[j,k]-gt[j,k])
      MSE=sumup/(width*height)
      return MSE
    



FLAG=True


images=[]
strokeMaps=[]
gts=[]
SharedMattes=[]

image=cv2.imread('/home/chen/Desktop/Comparing/image/GT01.png')
image=cv2.resize(image,(config.img_sz,config.img_sz))
images.append(image)
images=np.array(images)
strokeMaps.append(image)
strokeMaps=np.array(strokeMaps)
gt=cv2.imread('/home/chen/Desktop/Comparing/GT/GT01.png')
gt=gt[:,:,0]
gt=cv2.resize(gt,(config.img_sz,config.img_sz))
gts.append(gt)
gts=np.array(gts)
SharedMatte=cv2.imread('/home/chen/Desktop/Comparing/SharedMatting/GT01.png')
SharedMatte=cv2.resize(SharedMatte,(config.img_sz,config.img_sz))
SharedMattes.append(SharedMatte)
SharedMattes=np.array(SharedMattes)


time_step=0
MSEs=[]


while(FLAG and time_step<200):
  
  loc = np.random.uniform(-1.0,1.0,(1,2))
  strokeMaps = draw.drawstroke(strokeMaps,loc,strokeMaps,gts)
  alphas = cnn2.op(images,strokeMaps)
  matte=np.uint8(alphas[0])
  #cv2.imshow('After',matte)
  #cv2.waitKey(0)
  #cv2.destroyAllWindows()
  MSEs.append(MSE(matte,gt))
  #if MSE(matte,gt)<MSE(SharedMatte,gt):
  if MSE(matte,gt)<=36 and time_step>20:
     FLAG=False
  time_step+=1

print 'Time_Step is: ', time_step
print 'MSE of SharedMatte is: ', MSE(SharedMatte,gt)
print 'MSE of Random Method is: ', MSE(matte,gt)
print MSEs
